import os
import requests
from loguru import logger
import time

# 设置日志输出到文件
log_file = "upload_log.txt"
logger.add(log_file, rotation="10 MB", retention="7 days", level="INFO")

# 替换为你的接口URL
# api_url = "http://47.113.120.248:6010/knowledge_base/upload_docs"
api_url = "http://127.0.0.1:6010/knowledge_base/upload_docs"

uploaded_files_file = "name.txt"
knowledge_base_zh_name = "Mars"  # 替换为你的知识库中文名
# 替换为包含要发送的文件的文件夹路径
file_folder = fr"C:\Users\Adminis\Desktop\Mars知识库内容"

knowledgebase_name_zh2en = {
    "Mars": "Mars",
}
# 替换为你的参数值
params = {
    "knowledge_base_name": knowledgebase_name_zh2en[knowledge_base_zh_name],
    "override": True,  # 是否覆盖已有文件
    "to_vector_store": True,  # 是否上传文件后进行向量化
    "chunk_size": 500,  # 知识库中单段文本最大长度
    "chunk_overlap": 60,  # 知识库中相邻文本重合长度
    "zh_title_enhance": False,  # 是否开启中文标题加强
    "not_refresh_vs_cache": False,  # 暂不保存向量库
    # "add_title": False  # 是否在分段开头拼接文件名
}

file_list = os.listdir(file_folder)
file_list = sorted(file_list, key=lambda x: os.path.getsize(os.path.join(file_folder, x)))
start_time = time.time()

for file_name in file_list:
    file_path = os.path.join(file_folder, file_name)

    # 记录开始上传日志
    logger.info(f"开始上传: {file_name}")

    # 构建HTTP POST请求
    with open(file_path, 'rb') as file:
        files = {'files': (file_name, file)}
        try:
            response = requests.post(api_url, files=files, data=params, timeout=1200)
            if response.status_code == 200:
                logger.info(f"上传成功: {file_name}")
                with open(uploaded_files_file, 'a', encoding='utf-8') as txt_file:
                    txt_file.write(file_name + "\n")
            else:
                logger.error(f"上传失败 ----- {file_name}, 错误码 ----- {response.status_code}")
        except requests.exceptions.Timeout:
            logger.error(f"上传超时 ----- {file_name} ")
        except requests.exceptions.RequestException as e:
            logger.error(f"上传异常 ----- {file_name}: {e}")

end_time = time.time()
logger.info(f"程序结束, 总耗时 {end_time - start_time} 秒")
